#!/bin/env python
import pandas as pd
import numpy as np
import json
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import matplotlib.cm as cm

#----------
# Parameters
#----------
indir = "./Data"

ddomain = [123,146,24,46] 

tcols={
    "TD":"gray",
    "TS":"black",
    "C1":"blue",
    "C2":"green",
    "C3":"orange",
    "C4":"purple",
    "C5":"red",
    }

tsymbols={
    "TD":"D",
    "TS":"0",
    "C1":"1",
    "C2":"2",
    "C3":"3",
    "C4":"4",
    "C5":"5",
    }   

#----------
# Read Data
#----------

#--Precip data
infl_precip = "%s/data_composite_precip.json" % (indir)
df_ = pd.read_json(infl_precip)
lons = df_.lon.unique()
lats = df_.lat.unique()
flon,flat = np.meshgrid(lons,lats) # meshgrid
jmax,imax = np.shape(flon)
precip = df_['rain'].to_numpy().reshape((jmax,imax))

#--Tropical Cyclone data
infl_tc_info = "%s/data_composite_tc_info.json" % (indir)
with open(infl_tc_info) as json_fl:
    tcinfo = json.load(json_fl)
    
infl_tc_track = "%s/data_composite_tc_track.json" % (indir)
df_ = pd.read_json(infl_tc_track)
tlons = df_['lon'].to_numpy()
tlats = df_['lat'].to_numpy()
tcats = df_['cate'].to_numpy()

#--Front data
tsyear=1977
teyear=2015
INF="%s/data_composite_front_track.json" % (indir)
df_ = pd.read_json(INF)
qsf_ = df_.loc[(df_["dist"]>500) & (df_['iflag']==1)]
cf_ = df_.loc[(df_["dist"]>500) & (df_['iflag']==2)]
wf_ = df_.loc[(df_["dist"]>500) & (df_['iflag']==3)]

#----------
# Plot
#----------
fig = plt.figure(figsize=(10,13)) # set figure environemnt

ax = fig.add_subplot(1, 1, 1) # set up panel plot
m = Basemap(projection='mill',llcrnrlat=ddomain[2], urcrnrlat=ddomain[3], llcrnrlon=ddomain[0],urcrnrlon=ddomain[1], 
             lat_ts=10, resolution='h') # draw basemap
m.drawcoastlines() # draw coast line

meridians = np.arange(0.,360.,5.)
m.drawmeridians(meridians,labels=[0,0,0,1],linewidth=0, fontsize=14)
parallels = np.arange(0.,90, 5.)
m.drawparallels(parallels,labels=[1,0,0,1],linewidth=0, fontsize=14)

#--TC points
for tlon,tlat,tcat in zip(tlons,tlats,tcats):
    xx2,yy2 = m(tlon,tlat)
    m.plot(xx2,yy2, marker=r'$%s$' % (tsymbols[tcat]), color=tcols[tcat], ms=10, lw=10)
    
#--front 
#--quasi-stationary front
qlons = qsf_['lon'].to_numpy()
qlats = qsf_['lat'].to_numpy()
for lons, lats in zip(qlons,qlats):
    xx2,yy2 = m(lons,lats)
    m.plot(xx2, yy2, color='black', lw=2, alpha=0.1)


#--precip contour
contours=np.arange(0, 30.0, 2) # set contours

cmap=cm.Blues # set colormap
x,y=m(flon,flat) # lon,lat => x,y
cs = m.contourf(x,y,precip, contours, cmap=cmap, extend="both") # plot contours

#--title
plt.title("(a) Composite of Anomalous Days (1977\u20142015, May\u2014October)", fontsize=18)

textstr = '\n'.join((
   r'Total: %(ddays)i days' % tcinfo,
   r'TC-day Ratio: %(tcer).1f %%' % tcinfo,
   r'Mean storm intensity: %(meanws).1f $\rm m\ s^{-1}$' % tcinfo,
   r'Mean lifetime max intensity: %(maxws).1f $\rm m\ s^{-1}$' % tcinfo,
))   
props = dict(boxstyle='square', facecolor='white', alpha=1.0)
ax.text(0.025,0.975, textstr, transform=ax.transAxes, fontsize=16, horizontalalignment='left',verticalalignment='top', bbox=props)


# add legend
cax = fig.add_axes([0.1,  0.06, 0.85, 0.02]) 

art = plt.colorbar(cs, cax, orientation='horizontal')
art.set_label(r'Anomaly of 5-day mean daily precipitation [$\rm mm\ day^{-1}$]', fontsize=14)
art.ax.tick_params(labelsize=14)  


fig.tight_layout()
#plt.show()

fig.savefig("./Fig3a.png")

